import tensorflow as tf
import pandas as pd
import numpy as np
import nibabel as nib
from PIL import Image


def read_nii(filepath):
    """
    Reads .nii file and returns pixel array
    """
    ct_scan = nib.load(filepath)
    array = ct_scan.get_fdata()
    array = np.rot90(np.array(array))
    return array


class DataGen(tf.keras.utils.Sequence):
    def __init__(self, root, mask_name, transform,Img_size=128, batch_size=64):
        super(DataGen, self).__init__()
        raw_data = pd.read_csv(root + '/metadata.csv')
        for cc in raw_data.columns:
            raw_data[cc] = raw_data[cc].apply(lambda x: root + '/' + '/'.join(x.split('/')[-2:]))

        # Read sample
        self.sample_ct = []
        self.GT = []
        number = raw_data.shape[0]
        number = 2
        for i in range(number):
            """取每个CT图片层"""
            ct = read_nii(raw_data.loc[i, 'ct_scan']).astype(np.uint8)
            mask = read_nii(raw_data.loc[i, mask_name]).astype(np.uint8)
            num = ct.shape[-1]
            for j in range(num):
                im = Image.fromarray(ct[:, :, i]).resize(size=[Img_size, Img_size], resample=Image.BILINEAR)
                im = np.array(im).astype(np.float16)
                mk = Image.fromarray(mask[:, :, i]).resize(size=[Img_size, Img_size], resample=Image.BILINEAR)
                mk = np.array(mk)
                self.sample_ct.append(im[..., np.newaxis])
                self.GT.append(mk[..., np.newaxis])
            print('loading {}/{}'.format(i, number))
        self.sample_ct = np.array(self.sample_ct)
        print(self.sample_ct.shape)
        self.GT = np.array(self.GT)

        # normalize
        mins = self.sample_ct.min(axis=(1, 2, 3), keepdims=True)
        maxs = self.sample_ct.max(axis=(1, 2, 3), keepdims=True)
        self.sample_ct = (self.sample_ct - mins) / (maxs - mins)

        self.batch_size = batch_size
        self.transform = transform
        self.on_epoch_end()

    def __len__(self):
        return self.sample_ct.shape[0] // self.batch_size + 1

    def __getitem__(self, index):
        ct = self.sample_ct[index * self.batch_size:(index + 1) * self.batch_size]
        gt = self.GT[index * self.batch_size:(index + 1) * self.batch_size]
        if self.transform:
            ct = self.transform(ct)
            gt = self.transform(gt)
        return ct, gt

    def on_epoch_end(self):
        index = np.arange(self.sample_ct.shape[0])
        np.random.shuffle(index)
        self.sample_ct = self.sample_ct[index]
        self.GT = self.GT[index]


if __name__ == '__main__':
    root = 'E:\\2022Project\\data\\COVID-19-CT'
    G = DataGen(root, 'lung_mask', None, batch_size=10)
    for i in G:
        x, y = i
        print(x.shape, y.shape)
